#!/usr/bin/env python

#Original Script by RiCON; 
#http://forum.doom9.org/showthread.php?t=154535

#base on vfr.py 0.6
#modify by cunhan
#v0.2
#add: Adjust assfile extracted from ts according to trims
#v0.1
#add: in order to output aac file, using besplit to split audio

import sys
import re
import linecache
import optparse
import os
import random
import datetime
from subprocess import call

rat = re.compile('(\d+)(?:/|:)(\d+)')
v1re = re.compile('# timecode format v1')
v2re = re.compile('# timecode format v2')
trimre = re.compile("(?<!#)Trim\((\d+)\s*,\s*(\d+)\)",re.I)
tcConv = 'tcConv'
mkvmerge = 'mkvmerge'
beSplit = 'beSplit'

def main():
    
    linecache.clearcache()
    
    p = optparse.OptionParser(description='Grabs avisynth trims and outputs chapter file, qpfile and/or cuts audio (works with cfr and vfr input)',
                              prog='vfr.py',
                              version='VFR Chapter Creator Mod 0.1',
                              usage='%prog [options] infile.avs')
    p.add_option('--input', '-i', action="store", help='Audio file to be cut', dest="input")
    p.add_option('--output', '-o', action="store", help='Cut audio from MKVMerge', dest="output")
    p.add_option('--fps', '-f', action="store", help='Frames per second (for cfr input)', dest="fps")
    p.add_option('--timecodes', '-t', action="store", help='Timecodes file from the vfr video (v1 needs tcConv)', dest="timecodes")
    p.add_option('--chapters', '-c', action="store", help='Chapters file', dest="chapters")
    p.add_option('--qpfile', '-q', action="store", help='QPFile for x264', dest="qpfile")
    p.add_option('--verbose', '-v', action="store_true", help='Verbose', dest="verbose")
    p.add_option('--merge', '-m', action="store_true", help='Merge cut files', dest="merge")
    p.add_option('--remove', '-r', action="store_true", help='Remove cut files', dest="remove")
    p.add_option('--frames', action="store", help='Number of frames for v1 conversion', dest="frames")
    p.add_option('--test', action="store_true", help="Test mode (do not create new files)", dest="test")
    p.add_option('--beSplit', '-b', action="store_true", help='using BeSplit to Cut Audio file', dest="beSplit")
    p.add_option('--ass', action="store", help='Ajust ASS timestamp by trims', dest="ass")
    (options, args) = p.parse_args()
    
    if len(args) != 1:
        p.error("No avisynth script specified.")
    elif options.timecodes == None and os.path.isfile(args[0] + ".tc.txt") == False and options.fps == None:
        options.timecodes = '30000/1001'
    elif options.timecodes == None and os.path.isfile(args[0] + ".tc.txt") == True:
        options.timecodes = args[0] + ".tc.txt"
    elif options.timecodes != None and options.fps != None:
        p.error("Can't use vfr input AND cfr input")
    elif options.timecodes != None and os.path.isfile(options.timecodes) == True:
        options.timecodes = options.timecodes
    else:
        options.timecodes = options.fps
    
    if options.chapters != None:
        chapType = 'MKV' if options.chapters[-3:] == 'xml' else 'OGM'
    else:
        chapType = ''
    
    if options.output == None and options.input != None:
        options.output = '%s.%s.mka' % (options.input, args[0][:-4])
    
    if options.verbose == True:
        status = """
Avisynth file:   {input}
Audio file:      {audio}
Cut Audio file:  {cutaudio}
Timecodes/FPS:   {timecodes}
Chapters file:   {chapters}
QP file:         {qpfile}

Merge/Rem files: {merge}/{remove}
Verbose:         {verbose}
Test Mode:       {test}
""".format(input=args[0],
            audio=options.input,
            cutaudio=options.output,
            timecodes=options.timecodes,
            chapters=options.chapters,
            qpfile=options.qpfile,
            merge=options.merge,
            remove=options.remove,
            verbose=options.verbose,
            test=options.test)
        print(status)
    
    quiet = '' if options.verbose == True else '-q'
    audio = []
    Trims = []
    
    with open(args[0], "r") as avs:
        # use only the first non-commented line with trims
        for line in avs:
            if trimre.match(line) != None:
                Trims = trimre.findall(line)
                break
            
        if len(Trims) < 1:
            sys.exit("Error: Avisynth script has no uncommented trims")
        
        if options.verbose == True:
            print('In trims:  %s'%(', '.join(['(%s,%s)'%(i[0],i[1]) for i in Trims])))
        
        # trims' offset calculation
        Trims2 = []
        Trims2ts = []
        options.timecodes = [options.timecodes, determineFormat(options.timecodes)]
        if options.timecodes[1] == 2:
            tc = options.timecodes
            nTrims = int(options.frames) if options.frames != None else int(Trims[-1][1])+2
            if os.path.isfile(tc[0]+"v2.txt") == False:
                tcConv = call('"%s" "%s" "%s" %d' % (tcConv, tc[0], tc[0]+"v2.txt", nTrims))
                if tcConv > 0:
                    sys.exit("Failed to execute tcConv: %d; Please put it in your path" % tcConv)
            options.timecodes[0] = tc[0]+"v2.txt"
        for i in range(len(Trims)):
        
            fn1 = int(Trims[i][0])                      # first frame
            fn1ts = Ts(fn1,options.timecodes)         # first frame timecode
            fn2 = int(Trims[i][1])                      # last frame
            fn2ts = Ts(fn2,options.timecodes)         # last frame timecode
            fn2tsaud = Ts(fn2+1,options.timecodes)    # last frame timecode for audio
                        
            if i != 0:                                  # if it's not the first trim
                last = int(Trims[i-1][1])+1
                lastts = Ts(last,options.timecodes)
                offset += fn1-last
                offsetts += fn1ts-lastts
            elif fn1 > 0:                               # if the first trim doesn't start at 0
                offset = fn1
                offsetts = fn1ts
            else:
                offset = 0
                offsetts = 0
            
            # apply the offset to the trims
            Trims2.append([fn1-offset,fn2-offset])
            Trims2ts.append([fn1ts-offsetts,fn2ts-offsetts])
            
            # make list with timecodes to cut audio
            audio.append(fn1ts)
            audio.append(fn2tsaud)
    
    if options.verbose == True:
        print('Out trims: %s'%(', '.join(['(%d,%d)'%(i[0],i[1]) for i in Trims2])))
    
    # make qpfile
    if options.qpfile != None and options.test == None:
        with open(options.qpfile, "w") as qpf:
            for trim in Trims2:
                qpf.write('%s I -1\n' % trim[0])
    
    # adjust assfile extracted from ts according to trims
    if options.ass != None :
        #print audio
        if os.path.isfile(options.ass) == True :
            newasslines = []
            assfile = open(options.ass,"r")
            asslines = assfile.readlines()
            assfile.close()
            for line in asslines:
                if line.upper().startswith("DIALOGUE") == True :
                    tokens = line.split(",")
                    tokens[1] = adjustTimeStamp(tokens[1],audio)
                    tokens[2] = adjustTimeStamp(tokens[2],audio)
                    newline = ",".join(tokens)
                    newasslines.append(newline)
                else:
                    newasslines.append(line)
            newass = open(options.ass[:-4]+".adjust.ass","w")
            for line in newasslines:
                newass.write(line)
            newass.close()              
            
    # make audio cuts
    if options.input != None:
        delay = re.search('DELAY ([-]?\d+)',options.input).group(1) if re.search('DELAY ([-]?\d+)',options.input) != None else '0'
        if audio[0] == 0:
            includefirst = True
            audio = audio[1:]
        else:
            includefirst = False
        #using besplit
        if options.beSplit != None:
            cuttimes = ' '.join([str(a*1.0/1000000000) for a in audio])
            cutCmd = '"%s" -core( -input "%s" -prefix "%s" -type aac -a ) -split( %s )'%\
            (beSplit, options.input, options.output + '.split', cuttimes )
            if options.verbose == True:
                print 'Cutting: %s\n' % cutCmd
            if options.test == None:
                cutExec = call(cutCmd)
                #when beSplit Complete,it retrun 58 but no 0. that's a problem
                if cutExec > 0 and cutExec != 58:
                    sys.exit("Failed to execute beSplit: %d" % cutExec)
            if options.merge == True:
                merge = []
                for i in range(1,len(audio)):
                    if (i % 2 != 0):
                        merge.append('%s.split%02d.aac' % (options.output, i))
                listfile = os.path.join(os.path.split(options.input)[0],"join.txt")
                fh = open(listfile, "w")
                for l in merge:
                    fh.write(l + "\n")
                fh.close()
                mergeCmd = '"%s" -core( -input "%s" -prefix "%s" -join )' % \
                (beSplit,listfile, options.output)
                if options.verbose == True:
                    print 'Merging: %s\n' % ' '.join(merge)
                    print mergeCmd
                if options.test == None:
                    mergeExec = call(mergeCmd)
                    #
                    if mergeExec > 0 and mergeExec != 58:                        
                        sys.exit("Failed to execute besplit: %d" % mergeExec)
            if options.remove == True:
                remove = ['%s.split%02d.aac' % (options.output, i) for i in range(1,len(audio))]
                if options.verbose == True:
                    print('Removing: %s\n' % ', '.join(remove))
                if options.test == None:
                    [os.unlink(i) for i in remove]
                
        else:
            audio = [formatTime(i) for i in audio]
            cuttimes = ','.join(audio)
            cutCmd = '"%s" -o "%s" "%s" --sync 0:%s --split timecodes:%s %s' % (mkvmerge, options.output + '.split.mka', options.input, delay, cuttimes, quiet)
            if options.verbose == True:
                print('Cutting: %s\n' % cutCmd)
            if options.test == None:
                cutExec = call(cutCmd)
                #it's strang that mergeExec return 1 but not 0, even if the process successed.
#                if cutExec > 0:
#                    sys.exit("Failed to execute mkvmerge: %d" % cutExec)
            if options.merge == True:
                merge = []
                for i in range(1,len(audio)+2):
                    if (includefirst == True and i % 2 != 0) or (includefirst == False and i % 2 == 0):
                        merge.append('+%s.split-%03d.mka' % (options.output, i))
                merge[0] = merge[0][1:]
                mergeCmd = '"%s" -o "%s" %s %s' % (mkvmerge,options.output, ' '.join(merge), quiet)
                if options.verbose == True:
                    print('Merging: %s\n' % ' '.join(merge))
                if options.test == None:
                    mergeExec = call(mergeCmd)
                    if mergeExec > 0:
                        sys.exit("Failed to execute mkvmerge: %d" % mergeExec)
            
            if options.remove == True:
                remove = ['%s.split-%03d.mka' % (options.output, i) for i in range(1,len(audio)+2)]
                if options.verbose == True:
                    print('Removing: %s\n' % ', '.join(remove))
                if options.test == None:
                    [os.unlink(i) for i in remove]
            
    EditionUID = random.randint(100000,1000000)
    matroskaXmlHeader = '<?xml version="1.0" encoding="UTF-8"?>\n<!-- <!DOCTYPE Tags SYSTEM "matroskatags.dtd"> -->\n<Chapters>'
    matroskaXmlEditionHeader = """
	<EditionEntry>
		<EditionFlagHidden>{0}</EditionFlagHidden>
		<EditionFlagDefault>{1}</EditionFlagDefault>
		<EditionFlagOrdered>{2}</EditionFlagOrdered>
		<EditionUID>{3}</EditionUID>
""".format(0,1,1,EditionUID)
    matroskaXmlEditionFooter = '	</EditionEntry>'
    matroskaXmlFooter = '\n</Chapters>'
    
    matroskaXmlTagsHeader = '<?xml version="1.0" encoding="UTF-8"?>\n<!-- <!DOCTYPE Tags SYSTEM "matroskatags.dtd"> -->\n<Tags>'
    matroskaXmlTagsEdition = """
	<Tag>
		<Targets>
			<EditionUID>{0}</EditionUID>
			<TargetTypeValue>50</TargetTypeValue>
		</Targets>

		<Simple>
			<Name>TITLE</Name>
			<String>{1}</String>
			<TagLanguage>{2}</TagLanguage>
			<DefaultLanguage>1</DefaultLanguage>
		</Simple>

	</Tag>""".format(EditionUID,"Default","eng")
    
    if chapType == 'OGM':
        matroskaXmlHeader = matroskaXmlEditionHeader = matroskaXmlEditionFooter = matroskaXmlFooter = ''
    
    if options.test == None and chapType != '':
        with open(options.chapters, "w") as output:
        
            output.write(matroskaXmlHeader)
            
            output.write(matroskaXmlEditionHeader)
            
            [output.write(generateChap(formatTime(Trims2ts[i][0]), formatTime(Trims2ts[i][1]),i+1,chapType)) for i in range(len(Trims2ts))]
            
            output.write(matroskaXmlEditionFooter)
            
            output.write(matroskaXmlFooter)
    elif chapType != '':
        print("Writing %s Chapters to %s"%(chapType,options.chapters))

def adjustTimeStamp(timestr, timeInterval):
    
    temp = []
    temp.extend(timeInterval)
    temp.append(0)
    t =datetime.datetime.strptime(timestr,"%H:%M:%S.%f")
    #print timestr, t
    timestamp = ((t.hour*3600 +t.minute*60 + t.second)*1000000 + t.microsecond)*1000
    #print timestr, t, timestamp
    offset = 0
    for i in range(len(temp)):
        if i%2 == 0:
            offset = offset + (temp[i]-temp[i-1])
        if timestamp <= temp[i]:
            adjustTime =  timestamp - offset
            strT = formatTime(adjustTime)
            strT = strT[0:strT.index('.')+3]
            return strT
    
    return "00:00:00.000"

def formatTime(ts):
    
    s = ts // 1000000000
    ms = ts % 1000000000
    m = s // 60
    s = s % 60
    h = m // 60
    m = m % 60
    
    return '{0:02d}:{1:02d}:{2:02d}.{3:09d}'.format(h, m, s, ms)

def determineFormat(timecodes):
    if rat.match(timecodes) != None or re.match('\d+',timecodes) != None:
        return 1
    elif v1re.match(linecache.getline(timecodes,1)):
        return 2
    elif v2re.match(linecache.getline(timecodes,1)):
        return 3
    else:
        return 0

def Ts(fn,tc):
    
    # CFR
    if tc[1] == 1:
        fps = rat.search(tc[0]).groups() if rat.search(tc[0]) else [re.search('(\d+)',tc[0]).group(0),'1']
        ts = round((1000 * fn * float(fps[1])) / float(fps[0]),6)
    # VFR
    elif tc[1] == 2 or tc[1] == 3:
        ts = linecache.getline(tc[0],fn+2)
    else:
        sys.exit("Couldn't get timestamps")
        
    
    return int(float(ts)*1000000)

def generateChap(start, end, chapter, type):
    matroskaXml = """
		<ChapterAtom>
			<ChapterTimeStart>{0}</ChapterTimeStart>
			<ChapterTimeEnd>{1}</ChapterTimeEnd>
			<ChapterDisplay>
				<ChapterString>Chapter {2:02d}</ChapterString>
				<ChapterLanguage>{3}</ChapterLanguage>
			</ChapterDisplay>
		</ChapterAtom>
"""[1:].format(start,end,chapter,"eng")
    ogmTxt = 'CHAPTER{0:02d}={1}\nCHAPTER{0:02d}NAME=Chapter {0:02d}\n'.format(chapter,start[:-6])
    if type == 'MKV':
        return matroskaXml
    elif type == 'OGM':
        return ogmTxt

if __name__ == '__main__':
    main()
